package main

import (
	"compress/gzip"
	"net/http"
	"strconv"
	"strings"
)

type gzWriter struct {
	http.ResponseWriter
	http.Flusher

	w     http.ResponseWriter
	level int

	headerWritten bool
	gzw           *gzip.Writer
}

func (gw *gzWriter) Header() http.Header {
	return gw.w.Header()
}

func (gw *gzWriter) WriteHeader(statusCode int) {
	if statusCode == http.StatusOK {
		// 根据类型和内容长度判断是否可以压缩。
		contentType := gw.w.Header().Get("Content-Type")
		contentLength, err := strconv.Atoi(gw.w.Header().Get("Content-Length"))
		if err != nil {
			contentLength = -1
		}
		contentEncoding := gw.w.Header().Get("Content-Encoding")
		transferEncoding := gw.w.Header().Get("Transfer-Encoding")

		// 如果响应是text类型和json类型、长度超过1k、编码格式未知或者使用未压缩的编码，那么启用压缩。
		if (strings.HasPrefix(contentType, "text/") || strings.HasPrefix(contentType, "application/json") || strings.HasPrefix(contentType, "application/javascript")) &&
			(contentLength > 1024 || (contentLength < 0 && transferEncoding != "chunked")) &&
			(contentEncoding == "" || contentEncoding == "identity") {
			if gzw, err := gzip.NewWriterLevel(gw.w, gw.level); err == nil {
				gw.gzw = gzw
				gw.w.Header().Del("Content-Length")
				gw.w.Header().Set("Content-Encoding", "gzip")
				gw.w.Header().Set("Transfer-Encoding", "chunked")
			}
		}
	}

	gw.w.WriteHeader(statusCode)
	gw.headerWritten = true
}

func (gw *gzWriter) Write(data []byte) (int, error) {
	if !gw.headerWritten {
		gw.WriteHeader(http.StatusOK)
	}

	if gw.gzw != nil {
		// 如果已启用压缩，那么将数据写入压缩流。
		n, err := gw.gzw.Write(data)
		if err == nil {
			// 立刻刷新刚写入的数据。
			gw.gzw.Flush()
		}
		return n, err
	}
	return gw.w.Write(data)
}

func (gw *gzWriter) Flush() {
	if gw.gzw != nil {
		gw.gzw.Flush()
	}

	if fw, ok := gw.w.(http.Flusher); ok {
		fw.Flush()
	}
}
